package com.lmk.dsa.struct.tree;

import com.lmk.dsa.struct.commons.BstNode;

/**
 * 二叉查找树基本实现
 * @param <K>
 * @param <V>
 * @author LaoMake
 * @since 1.0
 */
public class BinarySearchTree<K extends Comparable<? super K>, V> implements BST<K, V> {
    /** 树的容量 */
    protected Integer capacity;

    /** 根节点 */
    protected BstNode<K, V> root;

    public BinarySearchTree() {
        capacity = 0;
    }

    @Override
    public int size() {
        return this.capacity;
    }

    @Override
    public int height() {
        return height(root);
    }

    /**
     * 返回指定节点的高度
     * @param node
     * @return
     */
    private int height(BstNode node) {
        if(node == null){
            return 0;
        }

        int left = height(node.left);
        int right = height(node.right);
        return Math.max(left, right) + 1;
    }

    @Override
    public void add(K key, V value) {
        BstNode node = new BstNode(key, value);
        root = add(root, node);
        capacity++;
    }

    @Override
    public void add(BstNode<K, V> node) {
        root = add(root, new BstNode(node));
        capacity++;
    }

    /**
     * 在指定节点下添加新的数据节点
     * @param parent
     * @param child
     * @return
     */
    protected BstNode<K, V> add(BstNode<K, V> parent, BstNode<K, V> child){
        if(parent == null){
            return child;
        }

        int flag = child.compareTo(parent);
        if(flag > 0){
            parent.right = add(parent.right, child);
        }else if(flag < 0){
            parent.left = add(parent.left, child);
        }else{
            // KEY相同，更新Value的值
            parent.value = child.value;
        }

        // 更新当前父节点的高度
        parent.height = 1 + Math.max(nodeHeight(parent.left), nodeHeight(parent.right));

        return parent;
    }

    @Override
    public BstNode<K, V> root() {
        return root;
    }

    @Override
    public BstNode<K, V> minimum() {
        if(capacity == 0){
            return null;
        }else{
            return minimum(root);
        }
    }

    /**
     * 返回以node为根的树中最小节点
     * @param node
     * @return
     */
    protected BstNode minimum(BstNode node) {
        if(node.left == null)
            return node;
        return minimum(node.left);
    }

    @Override
    public BstNode<K, V> maximum() {
        if(capacity == 0){
            return null;
        }else{
            return maximum(root);
        }
    }

    /**
     * 返回以node为根的树中最大节点
     * @param node
     * @return
     */
    protected BstNode maximum(BstNode node){
        if(node.right == null)
            return node;
        return maximum(node.right);
    }

    @Override
    public boolean contains(K key) {
        return get(key) != null;
    }

    @Override
    public BstNode<K, V> get(K key) {
        BstNode node = get(root, key);
        return node == null ? null : new BstNode(node.key, node.value);
    }

    /**
     * 指定节点，搜索数据
     * @param node
     * @param key
     * @return
     */
    protected BstNode<K, V> get(BstNode<K, V> node, K key) {
        if(node == null){
            return null;
        }

        int flag = key.compareTo(node.key);
        if(flag == 0){
            return node;
        }else if (flag < 0){
            return get(node.left, key);
        }else{
            return get(node.right, key);
        }
    }

    @Override
    public BstNode<K, V> remove(K key) {
        BstNode node = get(root, key);
        if(node == null){
            return null;
        }else{
            root = remove(root, key);
            capacity--;
            return new BstNode(node.key, node.value);
        }
    }

    /**
     * 从某个节点开始，根据key搜索删除节点
     * @param node
     * @param key
     * @return
     */
    protected BstNode<K, V> remove(BstNode<K, V> node, K key) {
        if(node == null){
            return null;
        }

        BstNode retNode = null;
        int flag = key.compareTo(node.key);
        if(flag < 0){
            node.left = remove(node.left, key);
            retNode = node;
        }else if(flag > 0){
            node.right = remove(node.right, key);
            retNode = node;
        }else{
            if(node.left == null){
                // 该节点的左子树为空
                BstNode right = node.right;
                node.right = null;
                retNode = right;
            }else if(node.right == null){
                // 该节点的右子树为空
                BstNode left = node.left;
                node.left = null;
                retNode = left;
            }else{
                // 左右子树均不为空
                // 找到比待删除的节点大的最小节点，即右子树中的最小节点
                BstNode target = minimum(node.right);
                target.right = removeMin(node.right);
                target.left = node.left;
                node.left = node.right = null;
                retNode = target; // 将目标节点返回给上级
            }
        }

        return retNode;
    }

    @Override
    public BstNode<K, V> removeMin() {
        if(capacity == 0){
            return null;
        }else{
            BstNode node = minimum(root);
            root = removeMin(root);
            capacity--;
            return new BstNode(node.key, node.value);
        }
    }

    // 删除掉以node为根的二分搜索树中的最小节点
    // 返回删除节点后新的二分搜索树的根节点
    protected BstNode<K, V> removeMin(BstNode<K, V> node) {
        // 处理递归到底的情况。
        if(node.left == null){
            BstNode rightNode = node.right;
            node.right = null;
            return rightNode;
        }

        node.left = removeMin(node.left);
        return node;
    }


    @Override
    public BstNode<K, V> removeMax() {
        if(capacity == 0){
            return null;
        }else{
            BstNode node = maximum(root);
            root = removeMax(root);
            capacity--;
            return new BstNode(node.key, node.value);
        }
    }

    // 删除掉以node为根的二分搜索树中的最大节点
    // 返回删除节点后新的二分搜索树的根节点
    protected BstNode<K, V> removeMax(BstNode<K, V> node){
        if(node.right == null){
            BstNode leftNode = node.left;
            node.left = null;
            return leftNode;
        }

        node.right = removeMax(node.right);
        return node;
    }

    @Override
    public boolean isBalanced() {
        return isBalanced(root);
    }

    /**
     * 判断当前节点的子树是否平衡
     * 原理：
     *      1.空节点是平衡的
     *      2.节点的平衡因子大于1，则代表不平衡
     *      3. 如果当前节点平衡，则判断左右子树是否也都平衡
     * @param node
     * @return
     */
    private boolean isBalanced(BstNode<K, V> node) {
        if(node == null){
            return true;
        }

        int balanceFactor = getBalanceFactor(node);
        if(Math.abs(balanceFactor) > 1){
            return false;
        }

        return isBalanced(node.left) && isBalanced(node.right);
    }

    /**
     * 获取节点的平衡因子
     * @param node
     * @return
     */
    protected int getBalanceFactor(BstNode<K, V> node){
        if(node == null){
            return 0;
        }else{
            return nodeHeight(node.left) - nodeHeight(node.right);
        }
    }

    /**
     * 获取节点的高度
     * @param node
     * @return
     */
    protected int nodeHeight(BstNode<K, V> node){
        return node != null ? node.height : 0;
    }
}
